import pandas as pd

def novice_stable_recommendation(
    stock_attr_path,
    top_n=5,
    vol_threshold=4.87,  # 75%分位数
    turnover_min=7.22,   # 50%分位数
    max_per_industry=1
):
    # 1. 读取数据（仅需股票属性文件，已包含行业信息）
    stock_attr = pd.read_csv(stock_attr_path)
    
    # 确认列名（确保有industry列）
    # print("=== 股票属性文件列名 ===")
    # print(stock_attr.columns.tolist())
    
    # 2. 筛选稳健股票
    stable_candidates = stock_attr[
        (stock_attr['volatility'] <= vol_threshold) &
        (stock_attr['avg_turnover_billion'] >= turnover_min)
    ].copy()
    print(f"\n筛选结果：共{len(stable_candidates)}只股票符合条件")
    
    if len(stable_candidates) == 0:
        return "没有符合条件的股票，请放宽阈值"
    
    # 3. 检查并处理行业信息（直接使用股票属性文件中的industry列）
    # 移除没有行业信息的股票
    stable_with_industry = stable_candidates.dropna(subset=['industry'])
    print(f"有效股票（含行业信息）：{len(stable_with_industry)}只")
    
    if len(stable_with_industry) == 0:
        return "没有包含行业信息的股票，请检查数据"
    
    # 4. 按行业控制多样性
    stable_sorted = stable_with_industry.sort_values(
        by=['volatility', 'avg_turnover_billion'],
        ascending=[True, False]  # 波动率低优先，同波动率下成交额高优先
    )
    
    recommended = []
    industry_count = {}
    for _, row in stable_sorted.iterrows():
        industry = row['industry']  # 直接使用现有industry列
        if industry_count.get(industry, 0) < max_per_industry:
            recommended.append(row)
            industry_count[industry] = industry_count.get(industry, 0) + 1
        if len(recommended) >= top_n:
            break
    
    # 5. 整理结果
    result = pd.DataFrame(recommended)[
        ['stock_code', 'stock_name', 'industry', 'volatility', 'avg_turnover_billion']
    ]
    result = result.rename(columns={
        'stock_code': '股票代码',
        'stock_name': '股票名称',
        'industry': '所属行业',
        'volatility': '波动率',
        'avg_turnover_billion': '平均成交额（亿元）'
    })
    return result


if __name__ == "__main__":
    # 只需股票属性文件路径（已包含行业信息）
    stock_attr_path = r"D:\pycode\StockRecommend\Stock_Recommend\hs300_data\graph_stock_nodes.csv"
    
    recommendations = novice_stable_recommendation(
        stock_attr_path=stock_attr_path,
        top_n=5,
        vol_threshold=4.87,
        turnover_min=7.22,
        max_per_industry=1
    )
    
    print("\n===== 新手稳健推荐股票列表 =====")
    print(recommendations)
